import collections
import contextlib
import sys
import wave
import os
import webrtcvad

#读取wav判断能否继续
def read_wave(path):
    #以二进制方式堆区音频数据
    with contextlib.closing(wave.open(path, 'rb')) as wf:
        #获取声道数，为1时继续执行
        num_channels = wf.getnchannels()
        assert num_channels == 1
        #采样字节长度，为2时继续执行
        sample_width = wf.getsampwidth()
        assert sample_width == 2
        #获取采样率，当采样率为8000，16000，32000时继续执行
        sample_rate = wf.getframerate()
        assert sample_rate in (8000, 16000, 32000)
        pcm_data = wf.readframes(wf.getnframes())
        #返回wav的二进制数据与采样率
        return pcm_data, sample_rate


def write_wave(path, audio, sample_rate):
    with contextlib.closing(wave.open(path, 'wb')) as wf:
        #设置声道
        wf.setnchannels(1)
        # 设置采样字节长度
        wf.setsampwidth(2)
        # 设置采样率
        wf.setframerate(sample_rate)
        # 写入数据
        wf.writeframes(audio)


class Frame(object):
    """Represents a "frame" of audio data."""

    def __init__(self, bytes, timestamp, duration):
        self.bytes = bytes
        self.timestamp = timestamp
        self.duration = duration


def frame_generator(frame_duration_ms, audio, sample_rate):
    #采样率*秒数*声道数=量化大小
    n = int(sample_rate * (frame_duration_ms / 1000.0) * 2)
    offset = 0
    timestamp = 0.0
    #量化大小对应的音频时间
    duration = (float(n) / sample_rate) / 2.0
    #如果offset起始量化大小加上量化大小小于音频 -->用来切分音频
    while offset + n < len(audio):
        #创建多个Frame类，用来存数据   audio[offset:offset + n]切出来的那一段数据，timestamp这段音频开始时间，duration这段音频的长度
        yield Frame(audio[offset:offset + n], timestamp, duration)
        timestamp += duration
        offset += n


def vad_collector(sample_rate, buffer_size, vad, frames):
    # 创建一个有两端的容器，数据溢出用来排除非人声
    ring_buffer = collections.deque(maxlen=buffer_size)

    triggered = False

    voiced_frames = []
    for frame in frames:
        # 把二进制数据和采样率交给vad
        is_speech = vad.is_speech(frame.bytes, sample_rate)
        #检测说话的开始
        if not triggered:
            # 在容器里添加一个对象（Frame类的实例，静音检测结果(True表示有人说)）
            ring_buffer.append((frame, is_speech))
            #f = Frame() speech = is_speech 得到有人声音的Frame的个数
            num_voiced = len([f for f, speech in ring_buffer if speech])
            # 当说话段数量大于缓冲区的90%时认为人声开始，所以进入if时前10段有人声恰好在ring_buffer里
            if num_voiced > 0.9 * ring_buffer.maxlen:
                # 通知循环已经找到了音频开始位置
                triggered = True
                #输出当前片段的开始时间
                # 当有9条是人声时，那么最先放进去的一条就是起始位置
                sys.stdout.write('+(%s)' % (ring_buffer[0][0].timestamp,))
                # 把音频的二进制数据放到列表
                for f, s in ring_buffer:
                    voiced_frames.append(f)
                #清空缓存区
                ring_buffer.clear()
        # 检测说话的结束
        else:
            # 上面清空了缓存区，所以现在开始写入有人声的数据
            voiced_frames.append(frame)
            # 把人声的数据写到列表
            ring_buffer.append((frame, is_speech))
            # 获取有多少静音段被写到了ring_buffer
            num_unvoiced = len([f for f, speech in ring_buffer if not speech])
            # 当静音段数量大于90%时执行
            if num_unvoiced > 0.9 * ring_buffer.maxlen:
                # 此时ring_buffer里有10段静音，那么第一段就是人声的结束位置
                sys.stdout.write('-(%s)' % (frame.timestamp + frame.duration))
                # 通知函数找到结束
                triggered = False
                # 暂停当前函数，并返回数据
                yield b''.join([f.bytes for f in voiced_frames])
                # 因为下次调用函数会从yield后面的语句开始所以删除上一段音频数据
                ring_buffer.clear()
                voiced_frames = []
    #如果到音频结束还没有找到结束时，执行到这里，把音频结束当做人声结束， 这里frame自然就是最后一个
    if triggered:
        sys.stdout.write('-(%s)' % (frame.timestamp + frame.duration))
    sys.stdout.write('\n')
    #列表部为空返回所有数据
    if voiced_frames:
        yield b''.join([f.bytes for f in voiced_frames])


def main():
    path = r"./"
    files = os.listdir(path)
    files = [path + "\\" + f for f in files if f.endswith('.wav')]
    con = 1
    for i in range(len(files)):
        args = [3, files[i]]
        #二进制数据，采样率
        audio, sample_rate = read_wave(args[1])
        # 创建VAD实例，并设置模式
        vad = webrtcvad.Vad(int(args[0]))
        #切分数据并返回一堆用来存放数据的Rrame类的生成器
        frames = frame_generator(30, audio, sample_rate)
        frames = list(frames)
        #获取人声段的二进制数据 10决定缓冲区大小
        segments = vad_collector(sample_rate, 10, vad, frames)
        # 给列表加上一个索引(多条音频时使用，segments被)
        for j, segment in enumerate(segments):
            path = './data/' + 'chunk-%002d.wav' % (con,)
            print(' Writing %s' % (path,))
            write_wave(path, segment, sample_rate)
            con += 1


if __name__ == '__main__':
    main()